import cv2
import os
import numpy as np
import sys
sys.path.append(os.path.abspath('.'))
from utils.util import sci_SSIM as SSIM 
from utils.util import SSIM_mean_param,SSIM_std_param,SSIM_cov_param
from .base_util import BaseUtil

def formula1(x,y,C1=1):
    ans = (2*x*y+C1)/(x*x + y*y +C1)
    return ans
# def SSIM_mean_param(ori_img,dst_img,C1=1):
#     mean_x = ori_img.mean()
#     mean_y = dst_img.mean()
#     ans = formula1(mean_x,mean_y)
#     return ans

# def SSIM_std_param(ori_img,dst_img,C1=1):
#     std_x = ori_img.std()
#     std_y = dst_img.std()
#     ans = formula1(std_x,std_y)
#     return ans

# def SSIM_cov_param(ori_img,dst_img,C1=1):
#     if(len(ori_img.shape) == 2):
#         h,w = ori_img.shape
#     std_x = ori_img.std()
#     std_y = dst_img.std()
#     cov = np.cov(np.stack((ori_img.reshape(h*w),dst_img.reshape(h*w))))[0,1]
#     ans = (cov + C1)/(std_x*std_y + C1)
#     return ans

def heat_map(ori_img,dst_img):
        # heat = ((dst_img - ori_img + 255)/2).astype(np.uint8)
        heat = np.abs(dst_img - ori_img).astype(np.uint8)
        return heat

def corr_heat_map(ori_img,dst_img,bh,bw):

    H,W = dst_img.shape
    corr_mat = np.zeros((H,W))
    
    for h in range(1,int(H/bh)+1):
        for w in range(1,int(W/bw)+1):

            fake_vec = dst_img[(h-1)*bh : h*bh, (w-1)*bw : w*bw  ].reshape(bh*bw)
            real_vec = ori_img[(h-1)*bh : h*bh, (w-1)*bw : w*bw  ].reshape(bh*bw)
            corr_mat[(h-1)*bh : h*bh , (w-1)*bw : w*bw] = np.corrcoef(fake_vec,real_vec)[0,1]
            pass
    # corr_mat = ((corr_mat+1)*255/2).astype(np.uint8)
    corr_mat = (np.abs(corr_mat)*255).astype(np.uint8)
    # heat_corr = cv2.applyColorMap(corr_mat,cv2.COLORMAP_JET)
    return corr_mat

def SSIM_heat_map(ori_img,dst_img,bh,bw):
    
    H,W = dst_img.shape
    SSIM_mat = np.zeros((H,W))
    
    for h in range(1,int(H/bh)+1):
        for w in range(1,int(W/bw)+1):
            fake_vec = dst_img[(h-1)*bh : h*bh, (w-1)*bw : w*bw  ]
            real_vec = ori_img[(h-1)*bh : h*bh, (w-1)*bw : w*bw  ]
            SSIM_mat[(h-1)*bh : h*bh , (w-1)*bw : w*bw] = SSIM(fake_vec,real_vec)

            pass
    SSIM_mat = (SSIM_mat*255).astype(np.uint8)
    # heat_ssim = cv2.applyColorMap(SSIM_mat,cv2.COLORMAP_JET)
    return SSIM_mat

def heatColor(img):
    return cv2.applyColorMap(img,cv2.COLORMAP_JET)

class CycleGANSABFFTUtil(BaseUtil):
    def __init__(self, opt):
        BaseUtil.__init__(self,opt)
        
    def phaseData(self):
        if not os.path.exists(self.opt.dir):
            return False
        self.opt.files = os.listdir(self.opt.dir)
        record = []
        self.img = {}
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if tmp[1] == 'fake_SAB.png':
                if not tmp[0] in record:
                    record.append(tmp[0])
                    fake = cv2.imread(os.path.join(self.opt.dir,file))
                    real = cv2.imread(os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SB.png'))
                    self.img[os.path.join(self.opt.dir,file)] = cv2.cvtColor(fake,cv2.COLOR_RGB2GRAY)
                    self.img[os.path.join(self.opt.dir,tmp[0]+'.png_s0_real_SB.png')] = cv2.cvtColor(real,cv2.COLOR_RGB2GRAY)
            
    def process(self):
        self.logfile = os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'uilt_log.txt')
        if os.path.exists(self.logfile):
            os.remove(self.logfile)
        for process in self.opt.processes:
            getattr(self,process)()
        # self.phaseData()
    
    def mean_heat_map(self):
        _dir = self.opt.dir
        record = []
        mean_heat_mat = np.zeros((256,256)).astype(np.float)
        count = 0
        print('Heat start')
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if tmp[1] == 'fake_SAB.png':
                if not tmp[0] in record:
                    record.append(tmp[0])
                    fake = self.img[os.path.join(_dir,file)]
                    real = self.img[os.path.join(_dir,tmp[0]+'.png_s0_real_SB.png')]
                    heat_mat = heat_map(real,fake)
                    mean_heat_mat += heat_mat
                    count += 1
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_heat_SAB.png'),heatColor(heat_mat))
        mean_heat_mat /= count
        self.heat_score = np.sum(mean_heat_mat) / (mean_heat_mat.shape[0]*mean_heat_mat.shape[1])
        x='None'
        with open(self.logfile,'a') as f:
            x = "heat Score:  " + str(self.heat_score)+'\n'
            f.write(x)
        print(x)
        # print("heat Score:  " + str(self.heat_score))
        mean_heat_mat = mean_heat_mat.astype(np.uint8)
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_heat.png'),heatColor(mean_heat_mat.astype(np.uint8)))
        return mean_heat_mat

    def mean_corr_heat_map(self):
        _dir = self.opt.dir
        bh=self.opt.bh
        bw=self.opt.bw
        
        # files = os.listdir(_dir)
        record = []
        mean_corr_mat = np.zeros((256,256)).astype(np.float)
        count = 0
        H,W = self.opt.imgh, self.opt.imgw
        print('corr start')
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if tmp[1] == 'fake_SAB.png':
                if not tmp[0] in record:
                    record.append(tmp[0])
                    #
                    # fake = cv2.imread(os.path.join(_dir,file))
                    # real = cv2.imread(os.path.join(_dir,tmp[0]+'.png_s0_real_SB.png'))
                    fake = self.img[os.path.join(_dir,file)]
                    real = self.img[os.path.join(_dir,tmp[0]+'.png_s0_real_SB.png')]
                    
                    # H,W = fake.shape

                    corr_mat = corr_heat_map(real,fake,bh,bw)
                    mean_corr_mat = (corr_mat + mean_corr_mat*count)/(count+1)
                    count += 1
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_hcorr_'+str(int(H/bh))+'X'+str(int(W/bw))+'_SAB.png'),heatColor(corr_mat))

                    pass
                
            # img = cv2.imread(os.path.join(_dir,file))
        mean_corr_mat = mean_corr_mat.astype(np.uint8)
        self.mean_corr_score = np.where(mean_corr_mat>255*0.6)[0].shape[0]/(H*W)
        # self.mean_corr_score = np.where(mean_corr_mat>mean_corr_mat.mean())[0].shape[0]/(H*W)
        x='None'
        with open(self.logfile,'a') as f:
            x = 'corr score:'+str(self.mean_corr_score)+'\n'
            f.write(x)
        print(x)
        # print('corr score:'+str(self.mean_corr_score))
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_hcorr_'+str(int(H/bh))+'X'+str(int(W/bw))+'.png'),heatColor(mean_corr_mat.astype(np.uint8)))
        return mean_corr_mat

    def mean_SSIM_heat_map(self):
        _dir = self.opt.dir
        bh=self.opt.bh
        bw=self.opt.bw
        
        # files = os.listdir(_dir)
                    
        record = []
        mean_SSIM_mat = np.zeros((256,256)).astype(np.float)
        count = 0
        H,W = self.opt.imgh, self.opt.imgw
        print('SSIM start')
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if tmp[1] == 'fake_SAB.png':
                if not tmp[0] in record:
                    record.append(tmp[0])
                    fake = self.img[os.path.join(_dir,file)]
                    real = self.img[os.path.join(_dir,tmp[0]+'.png_s0_real_SB.png')]
                    # H,W = fake.shape
                    SSIM_mat = SSIM_heat_map(real,fake,bh,bw)
                    
                    mean_SSIM_mat = (SSIM_mat + mean_SSIM_mat*count)/(count+1)
                    
                    count += 1
                    cv2.imwrite(os.path.join(_dir,tmp[0]+'.png_s0_SSIM_'+str(int(H/bh))+'X'+str(int(W/bw))+'_SAB.png'),heatColor(SSIM_mat))
                
                    pass
                
        # self.mean_SSIM_score = np.where(np.unique(mean_SSIM_mat)>255*0.7)[0].shape[0]/np.unique(mean_SSIM_mat).shape[0]?
        self.mean_SSIM_score = np.where(mean_SSIM_mat>255*0.6)[0].shape[0]/(H*W)
        # self.mean_SSIM_score = np.where(mean_SSIM_mat>mean_SSIM_mat.mean())[0].shape[0]/(H*W)
        x='None'
        with open(self.logfile,'a') as f:
            x = 'SSIM score:'+str(self.mean_SSIM_score)+'\n'
            f.write(x)
        print(x)
        # print('SSIM score:'+str(self.mean_SSIM_score))
        cv2.imwrite(os.path.join(self.opt.root,self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch,'mean_SSIM_'+str(int(H/bh))+'X'+str(int(W/bw))+'.png'),heatColor(mean_SSIM_mat.astype(np.uint8)))
        return mean_SSIM_mat

    def SSIM_tool(self):
        _dir = self.opt.dir
        record = []
        
        mean_p1 = 0
        mean_p2 = 0
        mean_p3 = 0
        mean_p123 = 0
        mean_ssim = 0
        
        count = 0
        print('SSIM tool start')
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if tmp[1] == 'fake_SAB.png':
                if not tmp[0] in record:
                    record.append(tmp[0])
                    fake = self.img[os.path.join(_dir,file)]
                    real = self.img[os.path.join(_dir,tmp[0]+'.png_s0_real_SB.png')]
                    # H,W = fake.shape
                    ssim = SSIM(fake,real)
                    p1 = SSIM_mean_param(real,fake)
                    p2 = SSIM_std_param(real,fake)
                    p3 = SSIM_cov_param(real,fake)

                    mean_p1 += p1
                    mean_p2 += p2
                    mean_p3 += p3
                    mean_p123 += p1*p2*p3
                    mean_ssim += ssim
                    
                    count+=1
                    pass
        mean_p1 /= count
        mean_p2 /= count
        mean_p3 /= count
        mean_p123 /= count
        mean_ssim /= count
        
        x='None'
        with open(self.logfile,'a') as f:
            x = "mean_SSIM: "+str(mean_ssim)+" ,p1: "+str(mean_p1)+" ,p2: "+str(mean_p2)+" ,p3: "+str(mean_p3)+" ,p123: "+str(mean_p123)+'\n'
            f.write(x)
        print(x)
        # print("mean_SSIM: "+str(mean_ssim)+" ,p1: "+str(mean_p1)+" ,p2: "+str(mean_p2)+" ,p3: "+str(mean_p3)+" ,p123: "+str(mean_p123))
        pass
    
    def FFID(self):
        _dir = self.opt.dir
        record = []
        from  numpy.fft import fft2,fftshift
        for file in self.opt.files:
            tmp = file.split('.png_s0_')
            if tmp[1] == 'fake_SAB.png':
                record.append(tmp[0])
                fake = self.img[os.path.join(_dir,file)]
                real = self.img[os.path.join(_dir,tmp[0]+'.png_s0_real_SB.png')]
                F_fake = fftshift(fft2(fake))
                F_real = fftshift(fft2(real))
                FFID = np.sum(F_fake-F_real)
        pass
    
    def FID(self):
        pass